package org.momo.rpc.ratelimit.impl;


import org.momo.rpc.ratelimit.RateLimiter;


public class TokenBucketRateLimiter implements RateLimiter {

    private final int capacity;

    private final int rate;

    private int tokens;

    private long lastRefillTimestamp;

    public TokenBucketRateLimiter(int capacity, int rate) {
        this.capacity = capacity;
        this.rate = rate;
        this.tokens = capacity;
        this.lastRefillTimestamp = System.currentTimeMillis();
    }

    @Override
    public synchronized boolean limit() {
        refill();
        if (tokens >= 1) {
            tokens -= 1;
            return true;
        }
        return false;
    }

    private void refill() {
        long now = System.currentTimeMillis();
        int newTokens = (int) ((now - lastRefillTimestamp) * rate / 1000);
        if (newTokens > 0) {
            tokens = Math.min(capacity, tokens + newTokens);
            lastRefillTimestamp = now;
        }
    }

}
